!pip install numpy
!pip install matplotlib>=3.5.1  # 3.5.1未満であるとグラフの表示が正常に行われない可能性あり
!pip install pandas
!pip install dwave-ocean-sdk
!pip install openjij

import math
from decimal import Decimal, ROUND_HALF_UP, ROUND_HALF_EVEN
import time

import numpy as np
import pandas as pd
import matplotlib.cm as cm
import matplotlib.pyplot as plt

%matplotlib inline

from dwave.system.samplers import DWaveSampler
from dwave.system.composites import EmbeddingComposite
from openjij import SASampler

N = 20
Avg = 10
Var = 10

SD = np.sqrt(Var)
C = np.random.normal(loc=Avg, scale=SD, size=N)
W = np.round(C)  # データセットを整数値に丸め込む

def plot_numbers(numbers):
    x_list = np.arange(len(numbers))
    plt.bar(x_list, numbers, color="skyblue")
    for i, n in enumerate(numbers):
        plt.text(i, n / 2, n, ha="center")
    plt.show()

plot_numbers(W)

# 生成したデータセットの合計値、平均値、分散を表示
print("合計:\t", sum(W))
print("平均:\t", sum(W) / N)
print("分散:\t", np.var(W))

t1 = time.perf_counter()

# 全通りのグループ分けの結果
Group_A_BF = np.zeros(((2**N), N))
Group_B_BF = np.zeros(((2**N), N))

# bitによる組み合わせ
Combo = np.zeros(((2**N), N))

# 2つのグループの差の最小値
Diff_BF = sum(W)

# 最適解のグループ分けの暫定値
note = 0

for i in range(2**N):
    # bitで組み合わせを表現するためにiを2進数に変換し、残りを0で埋める
    Bit = format(i, "b").zfill(N)

    for j in range(N):
        # Bitの0,1に応じて全通りの組合せを保存する
        if Bit[j] == "1":
            Group_A_BF[i][j] = W[j]
        else:
            Group_B_BF[i][j] = W[j]

    # 上で求めた全通りの組合せのうち、差が最小となるものを抽出する
    Diff_tmp = abs(sum(Group_A_BF[i]) - sum(Group_B_BF[i]))

    if Diff_tmp <= Diff_BF:
        Diff_BF = Diff_tmp
        note = i

t2 = time.perf_counter()

# 計算時間を記録する
Time_BF = t2 - t1

def plot_grouping(A, B, bar_width=0.4, alpha=1):
    df = pd.DataFrame([A, B], index=["A", "B"]).T
    df_nan = df.replace([0], np.nan)

    fig, ax = plt.subplots()

    Sum = np.array([sum(A), sum(B)])

    data_points = np.arange(len(df.columns))
    bottom_data = pd.Series(np.zeros(len(df.columns)), index=df.columns.tolist())

    index = np.array(["A", "B"])

    for i in range(len(df.index)):
        p1 = ax.bar(
            index,
            df.iloc[i],
            bar_width,
            bottom=bottom_data,
            alpha=alpha,
            color="skyblue",
            edgecolor="white",
            linewidth=0.5,
        )
        ax.bar_label(p1, label_type="center")

        bottom_data += df.iloc[i]
        ax.set_xticks(data_points)
        ax.set_xticklabels(df.columns)

    return plt.show()

plot_grouping(Group_A_BF[note], Group_B_BF[note])

print("合計値の差\t", Diff_BF)
print("計算時間[s]\t", Time_BF)

Q = {}
for i in range(N):
    Q[i, i] = 4 * W[i] * (W[i] - sum(W))
    for j in range(i + 1, N):
        Q[i, j] = 8 * W[i] * W[j]

def show_qubo(qubo, cmap=cm.GnBu, save_path=None):
    n_qubo = (
        max(sorted(qubo.keys())[-1][0], sorted(qubo.keys(), key=lambda x: x[1])[-1][1])
        + 1
    )

    np_qubo = np.zeros((n_qubo, n_qubo))
    for (pos_x, pos_y), coeff in qubo.items():
        np_qubo[pos_x][pos_y] = coeff

    plt.imshow(np_qubo, cmap=cmap)
    plt.colorbar()

    if save_path is not None:
        plt.savefig(save_path)

    plt.show()

show_qubo(Q)

# 接続情報
endpoint = "https://cloud.dwavesys.com/sapi"
token = "YOUR_TOKEN"
solver = "Advantage_system6.2"

dw = DWaveSampler(endpoint=endpoint, token=token, solver=solver)

# D-Waveマシンのハードウェアグラフに埋め込みを行う
sampler = EmbeddingComposite(dw)

# 一度に取得するサンプル数
num_reads = 300

# D-Waveマシンのアニーリング時間
annealing_time = 20

# D-Waveマシンによるサンプリングを実行する
response_DW = sampler.sample_qubo(Q, num_reads=num_reads, annealing_time=annealing_time)

response_DW.to_pandas_dataframe().head()

response_DW.info

# アニーリングマシンから返された結果のエネルギーを計算する関数
def calculate_energy(solution):
    Ene = sum(W) ** 2
    for i in range(N):
        Ene += Q[i, i] * solution[i]
        for j in range(i + 1, N):
            Ene += Q[i, j] * solution[i] * solution[j]
    return Ene

# 最適解と近似解の個数をカウントする関数
def count_solutions(response):
    # 最適解と近似解の個数
    num_optimal_sol = 0
    num_approx_sol = 0

    # 最適解と近似解
    optimal_sol = []
    approx_sol = []

    # エネルギー
    Energies = []

    for state in response.record:
        energy = calculate_energy(state.sample)

        # 出現回数分だけエネルギーを加えることで、後で全体のエネルギー分布を把握する
        Energies += [energy] * state.num_occurrences

    # ビット全探索の結果を元に、最適解を見つける
    for state in response.record:
        energy = calculate_energy(state.sample)

        if energy == Diff_BF:
            num_optimal_sol += state.num_occurrences
            optimal_sol.append(state.sample)

        if energy == min(Energies):
            num_approx_sol += state.num_occurrences
            approx_sol.append(state.sample)

    return num_optimal_sol, num_approx_sol, optimal_sol, approx_sol, Energies

num_optimal_sol, num_approx_sol, optimal_sol, approx_sol, Energies = count_solutions(
    response_DW
)

print("最適解（個）\t", num_optimal_sol)
print("近似解（個）\t", num_approx_sol)

Group_A_DW = []
Group_B_DW = []

for solution in approx_sol:
    Group_A = np.zeros(N)
    Group_B = np.zeros(N)

    for i in range(N):
        if solution[i] == 1:
            Group_A[i] = W[i]
        else:
            Group_B[i] = W[i]

    Group_A_DW.append(Group_A)
    Group_B_DW.append(Group_B)

# 得られた近似解の組合せを1つ表示する
S = 0  # 解のインデックス (0 <= S <= num_optimal_sol - 1)
plot_grouping(Group_A_DW[S], Group_B_DW[S])

print("2つのグループの合計値の差\t", abs(sum(Group_A_DW[S]) - sum(Group_B_DW[S])))
print("計算時間（理論値）[s]\t", annealing_time * num_reads / 10**6)
print("計算時間（実測値）[s]\t", response_DW.info["timing"]["qpu_sampling_time"] / 10**6)

# 解のエネルギー分布を可視化する関数
def plot_energies(energies):
    plt.hist(energies, color="skyblue")
    plt.xlabel("Energy")
    plt.ylabel("Frequency")
    plt.show()

plot_energies(Energies)

# TTSを計算する関数（Tau[s]）
def calculate_TTS(p_d, Tau, p_s):
    if p_s == 0:
        raise ValueError("p_s!=0 is required")
    return Tau * np.log(1 - p_d) / np.log(1 - p_s)

p_d = 0.99  # 0.99の場合、99%の確率で最適解が得られるまでの計算時間
Tau = annealing_time * 10 ** (-6)
p_s = num_optimal_sol / num_reads

print("TTS[µs]", calculate_TTS(p_d, Tau, p_s) * 10**6)

# SAをサンプラーとして選択
sampler = SASampler()

# 一度に取得するサンプル数
num_reads = 300

# アニーリング中のパラメータ(温度, 横磁場)を下げていくときの分割数
# この値を増やすほどゆっくりアニーリングをすることに相当し、アニーリング時間が伸びる
num_sweeps = 50

# SAによるサンプリングを実行する
response_SA = sampler.sample_qubo(Q, num_reads=num_reads, num_sweeps=num_sweeps)

response_SA.to_pandas_dataframe().head()

num_optimal_sol, num_approx_sol, optimal_sol, approx_sol, Energies = count_solutions(
    response_SA
)

print("最適解（個）\t", num_optimal_sol)
print("最適解（個）\t", num_approx_sol)

Group_A_SA = []
Group_B_SA = []

for solution in approx_sol:
    Group_A = np.zeros(N)
    Group_B = np.zeros(N)

    for i in range(N):
        if solution[i] == 1:
            Group_A[i] = W[i]
        else:
            Group_B[i] = W[i]

    Group_A_SA.append(Group_A)
    Group_B_SA.append(Group_B)

# 得られた近似解の組合せを1つ表示する
S = 0  # 解のインデックス (0 <= S <= num_optimal_sol - 1)
plot_grouping(Group_A_SA[S], Group_B_SA[S])

print("2つのグループの合計値の差\t", abs(sum(Group_A_SA[S]) - sum(Group_B_SA[S])))
print("計算時間[µs]\t", response_SA.info["execution_time"])

plot_energies(Energies)